const H0: u32 = 0x6a09e667;
const H1: u32 = 0xbb67ae85;
const H2: u32 = 0x3c6ef372;
const H3: u32 = 0xa54ff53a;
const H4: u32 = 0x510e527f;
const H5: u32 = 0x9b05688c;
const H6: u32 = 0x1f83d9ab;
const H7: u32 = 0x5be0cd19;

const K: [u32; 64] = [
   0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
   0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
   0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
   0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
   0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
   0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
   0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
   0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
];

type MessageDigest = [u32; 8];

pub fn compression(bytes: Vec<u8>) -> [u8; 32] {  
   assert_eq!(bytes.len() % 64, 0, "Padded message's byte size must a multiple of 64");
   
   // for each 512 bits(64 bytes) as a block
   let md = bytes
      .chunks(64).into_iter()
      .fold([H0, H1, H2, H3, H4, H5, H6, H7], calculate_new_md);
   md_to_256bits(&md)
}

#[allow(unused)]
fn calculate_new_md_less_w_space_pipline(md: MessageDigest, block: &[u8]) -> MessageDigest {
   assert_eq!(block.len(), 64, "Block's byte size must be 64");
   
   let mut w: [u32; 16] = [0; 16];

   let (mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h) = 
      (md[0], md[1], md[2], md[3], md[4], md[5], md[6], md[7]);


   // 这论的 W 必须先计算，才能进行后续计算。
   // 但可以把 W 与 a-h 并行计算（W 计算领先 a-h），这样每一周期，W 的值可以直接使用

   // 首先计算第一个 W
   let mut bytes: [u8; 4] = [0; 4];
   bytes.copy_from_slice(&block[0..4]);
   w[0] = u32::from_be_bytes(bytes);

   let mut wt_add_k_add_h: u32 = K[0] + w[0] + h;

   // 以下代码无法正常运行，但是解释了如何并行加速
   // 两个模块并行计算
   for i in 0..64 {
      // 模块 1：计算 wt
      let j = i + 1;
      if j < 63 {
         let wt: u32 = if j < 16 {
            let mut bytes: [u8; 4] = [0; 4];
            bytes.copy_from_slice(&block[(j*4)..(j*4 + 4)]);
            w[j] = u32::from_be_bytes(bytes);
            w[j]
         } else {
            let s0 = right_rotate(w[1], 7) ^ right_rotate(w[1], 18) ^ (w[1] >> 3);
            let s1 = right_rotate(w[14], 17) ^ right_rotate(w[14], 19)  ^ (w[14] >> 10);
            let mut wt: u32 = 0;
            wt = wt.wrapping_add(s0);
            wt = wt.wrapping_add(w[9]);
            wt = wt.wrapping_add(s1);
            
            for i in 0..15 {
               w[i] = w[i + 1]
            }
            w[15] = wt;
   
            wt
            
         };

         // 这论的 g 是下一轮的 h
         wt_add_k_add_h = wt_add_k_add_h + wt + g; 
      } 

      // 模块 2：计算 a-h
      let s1 = right_rotate(e, 6) ^ right_rotate(e, 11) ^ right_rotate(e, 25);
      let ch = (e & f) ^ ((!e) & g);
      let mut t1 = s1;
      t1 = t1.wrapping_add(ch);
      // 直接加 wt_add_k_add_h: 替代 w + k + h
      t1 = t1.wrapping_add(wt_add_k_add_h);
      let t1 = t1;

      let s0 = right_rotate(a, 2) ^ right_rotate(a, 13) ^ right_rotate(a, 22);
      let maj = (a & b) ^ (a & c) ^ (b & c);
      let t2 = s0.wrapping_add(maj);

      h = g;
      g = f;
      f = e;
      e = d.wrapping_add(t1);
      d = c;
      c = b;
      b = a;
      a = t1.wrapping_add(t2);
   }

   [
      md[0].wrapping_add(a),
      md[1].wrapping_add(b),
      md[2].wrapping_add(c),
      md[3].wrapping_add(d),
      md[4].wrapping_add(e),
      md[5].wrapping_add(f),
      md[6].wrapping_add(g),
      md[7].wrapping_add(h),
   ]
}

#[allow(unused)]
fn calculate_new_md_less_w_space(md: MessageDigest, block: &[u8]) -> MessageDigest {
   assert_eq!(block.len(), 64, "Block's byte size must be 64");

   let mut w: [u32; 16] = [0; 16];

   let (mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h) = 
      (md[0], md[1], md[2], md[3], md[4], md[5], md[6], md[7]);


   // 这论的 W 必须先计算，才能进行后续计算。
   // 但可以把 W 与 a-h 并行计算（W 计算领先 a-h），这样每一周期，W 的值可以直接使用
   for i in 0..64 {
      let wt: u32 = if i < 16 {
         let mut bytes: [u8; 4] = [0; 4];
         bytes.copy_from_slice(&block[(i*4)..(i*4 + 4)]);
         w[i] = u32::from_be_bytes(bytes);
         w[i]
      } else {
         let s0 = right_rotate(w[1], 7) ^ right_rotate(w[1], 18) ^ (w[1] >> 3);
         let s1 = right_rotate(w[14], 17) ^ right_rotate(w[14], 19)  ^ (w[14] >> 10);
         let mut wt: u32 = 0;
         wt = wt.wrapping_add(s0);
         wt = wt.wrapping_add(w[9]);
         wt = wt.wrapping_add(s1);
         
         for i in 0..15 {
            w[i] = w[i + 1]
         }
         w[15] = wt;

         wt
         
      };

      let s1 = right_rotate(e, 6) ^ right_rotate(e, 11) ^ right_rotate(e, 25);
      let ch = (e & f) ^ ((!e) & g);
      let mut t1 = h.wrapping_add(s1);
      t1 = t1.wrapping_add(ch);
      t1 = t1.wrapping_add(K[i]);
      t1 = t1.wrapping_add(wt);
      let t1 = t1;

      let s0 = right_rotate(a, 2) ^ right_rotate(a, 13) ^ right_rotate(a, 22);
      let maj = (a & b) ^ (a & c) ^ (b & c);
      let t2 = s0.wrapping_add(maj);

      h = g;
      g = f;
      f = e;
      e = d.wrapping_add(t1);
      d = c;
      c = b;
      b = a;
      a = t1.wrapping_add(t2);
   }

   [
      md[0].wrapping_add(a),
      md[1].wrapping_add(b),
      md[2].wrapping_add(c),
      md[3].wrapping_add(d),
      md[4].wrapping_add(e),
      md[5].wrapping_add(f),
      md[6].wrapping_add(g),
      md[7].wrapping_add(h),
   ]
}

fn calculate_new_md(md: MessageDigest, block: &[u8]) -> MessageDigest {
   assert_eq!(block.len(), 64, "Block's byte size must be 64");
   // For each block, loop 64 times to get new md
   // let mut block = block.iter().cloned().collect::<Vec<_>>();

   // 1. Calculate w[0..63]
   let mut w: [u32; 64] = [0; 64];
   for i in 0..16 {
      let mut bytes: [u8; 4] = [0; 4];
      bytes.copy_from_slice(&block[(i*4)..(i*4 + 4)]);
      w[i] = u32::from_be_bytes(bytes);
   }

   for i in 16..64 {
      let s0 = right_rotate(w[i - 15], 7) ^ right_rotate(w[i - 15], 18) ^ (w[i - 15] >> 3);
      let s1 = right_rotate(w[i - 2], 17) ^ right_rotate(w[i - 2], 19)  ^ (w[i - 2] >> 10);
      w[i] = w[i - 16].wrapping_add(s0);
      w[i] = w[i].wrapping_add(w[i - 7]);
      w[i] = w[i].wrapping_add(s1);
   }

   // 2. Init a to h
   let (mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h) = 
      (md[0], md[1], md[2], md[3], md[4], md[5], md[6], md[7]);

   // 3. Loop 64 times to update a to h
   for i in 0..64 {
      let s1 = right_rotate(e, 6) ^ right_rotate(e, 11) ^ right_rotate(e, 25);

      let ch = (e & f) ^ ((!e) & g);

      let mut t1 = h.wrapping_add(s1);
      t1 = t1.wrapping_add(ch);
      t1 = t1.wrapping_add(K[i]);
      t1 = t1.wrapping_add(w[i]);
      let t1 = t1;

      let s0 = right_rotate(a, 2) ^ right_rotate(a, 13) ^ right_rotate(a, 22);

      let maj = (a & b) ^ (a & c) ^ (b & c);

      let t2 = s0.wrapping_add(maj);
      
      h = g;
      g = f;
      f = e;
      e = d.wrapping_add(t1);
      d = c;
      c = b;
      b = a;
      a = t1.wrapping_add(t2);

   }

   // 4. return new md
   [
      md[0].wrapping_add(a),
      md[1].wrapping_add(b),
      md[2].wrapping_add(c),
      md[3].wrapping_add(d),
      md[4].wrapping_add(e),
      md[5].wrapping_add(f),
      md[6].wrapping_add(g),
      md[7].wrapping_add(h),
   ]
}

fn md_to_256bits(md: &MessageDigest) -> [u8; 32] {
   let mut res = [0; 32];
   for (u32_idx, m) in md.into_iter().enumerate() {
      for (byte_idx, byte) in m.to_be_bytes().into_iter().enumerate() {
         res[u32_idx * 4 + byte_idx] = byte;
      }
   }

   res
}

fn right_rotate(n: u32, d: usize) -> u32 {
   (n >> d) | (n << (32 - d))
}
